Skip to content

[Json] Use partition and take in RunEndEncoded decoder - #9658

Merged
Jefffrey merged 5 commits into
apache:mainfrom
liamzwbao:issue-9645-ree-optimization
Apr 19, 2026
Merged

[Json] Use partition and take in RunEndEncoded decoder#9658
Jefffrey merged 5 commits into
apache:mainfrom
liamzwbao:issue-9645-ree-optimization

Conversation

@liamzwbao

@liamzwbao liamzwbao commented Apr 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Optimize RunEndEncoded decoder to use partition and take, substantially improving performance (over 2x speedup).

Are these changes tested?

Covered by existing tests

Are there any user-facing changes?

No

@github-actions github-actions Bot added the arrow Changes to the arrow crate label Apr 3, 2026
@liamzwbao
liamzwbao force-pushed the issue-9645-ree-optimization branch from 568cceb to ffd5849 Compare April 3, 2026 21:46
@liamzwbao
liamzwbao marked this pull request as ready for review April 3, 2026 21:49
@liamzwbao

liamzwbao commented Apr 3, 2026

Copy link
Copy Markdown
Contributor Author

Local benchmark results

Benchmarking decode_short_ree_runs_json/131072: Collecting 50 samples in estimated 5.1127 s (650 ite
decode_short_ree_runs_json/131072
        time:   [7.7892 ms 7.8779 ms 7.9584 ms]
        thrpt:  [217.23 MiB/s 219.45 MiB/s 221.95 MiB/s]
 change:
        time:   [−49.865% −49.220% −48.647%] (p = 0.00 < 0.05)
        thrpt:  [+94.730% +96.929% +99.463%]
        Performance has improved.

Benchmarking decode_long_ree_runs_json/131072: Collecting 50 samples in estimated 5.1862 s (700 iter
decode_long_ree_runs_json/131072
        time:   [6.5677 ms 6.6715 ms 6.7831 ms]
        thrpt:  [223.96 MiB/s 227.71 MiB/s 231.31 MiB/s]
 change:
        time:   [−56.851% −56.093% −55.362%] (p = 0.00 < 0.05)
        thrpt:  [+124.03% +127.75% +131.75%]
        Performance has improved.

Benchmarking decode_short_ree_runs_serialize: Collecting 100 samples in estimated 5.8899 s (600 iter
decode_short_ree_runs_serialize
        time:   [8.9996 ms 9.0555 ms 9.1159 ms]
        change: [−49.254% −48.848% −48.441%] (p = 0.00 < 0.05)
        Performance has improved.

Benchmarking decode_long_ree_runs_serialize: Collecting 100 samples in estimated 5.2928 s (600 itera
decode_long_ree_runs_serialize
        time:   [8.4019 ms 8.4652 ms 8.5338 ms]
        change: [−52.747% −52.330% −51.944%] (p = 0.00 < 0.05)
        Performance has improved.

alamb
alamb approved these changes Apr 16, 2026
View reviewed changes

@alamb alamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me -- thank you @liamzwbao

Comment thread arrow-json/Cargo.toml
arrow-buffer = { workspace = true }
arrow-cast = { workspace = true }
arrow-data = { workspace = true }
arrow-ord = { workspace = true }

@alamb alamb Apr 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are fairly non trivial crates (ord and select) so it is sad to see the dependencies being added here

That being said, I think it is becoming clear that anything involving REE benefits from those two

Maybe we could split them up or osmething into new crates with the "core" parts (specifically partition and take) 🤔

arrow-take and arrow-partition maybe 🤔

I'll file a ticket to consider this

@alamb alamb Apr 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Jefffrey Jefffrey Apr 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth noting that arrow-cast already brings in ord & select in it's dependencies, so I don't think this actually affects the build too much (compared to previous state)

@alamb alamb Apr 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth noting that arrow-cast already brings in ord & select in it's dependencies, so I don't think this actually affects the build too much (compared to previous state)

yeah, adding REE to cast required those new dependencies. We have some ideas of how to avoid it but they aren't great

alamb
alamb reviewed Apr 16, 2026
View reviewed changes
Comment thread arrow-json/src/reader/run_end_array.rs
let values_data = mutable.freeze();
let run_ends_data =
PrimitiveArray::<R>::new(ScalarBuffer::from(run_ends), None).into_data();
let indices = UInt32Array::from_iter_values(indices.into_iter().map(|i| i as u32));

@alamb alamb Apr 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In theory the old code could also handle usize indices not ust u32 but I think in practice it won't matter

@liamzwbao liamzwbao Apr 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, also indices are bounded by the tape pos (which is &[u32]), so they should never exceed u32::MAX IIUC

Jefffrey
Jefffrey approved these changes Apr 19, 2026
View reviewed changes

@Jefffrey Jefffrey left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This behaviour of efficiently generating a REE from a flat array seems generic enough to be worth splitting into a separate function (that maybe other users might need), but could be a followup 🤔

@Jefffrey
Jefffrey merged commit fa1cc58 into apache:main Apr 19, 2026
24 checks passed
@alamb

alamb commented Apr 20, 2026

Copy link
Copy Markdown
Contributor

This behaviour of efficiently generating a REE from a flat array seems generic enough to be worth splitting into a separate function (that maybe other users might need), but could be a followup 🤔

It also exists in the cast kernels -- so refactoring that out (or reusing the cast) would be good

@alamb alamb mentioned this pull request Apr 23, 2026
Closed
@liamzwbao
liamzwbao deleted the issue-9645-ree-optimization branch April 26, 2026 15:13
Rich-T-kid pushed a commit to Rich-T-kid/arrow-rs that referenced this pull request Jun 2, 2026
@liamzwbao @Rich-T-kid
[Json] Use partition and take in RunEndEncoded decoder (apache#9658)
ea11e7c 
# Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->

- Closes apache#9645 .

# Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->

# What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

Optimize RunEndEncoded decoder to use `partition` and `take`,
substantially improving performance (over 2x speedup).

# Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

Covered by existing tests

# Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.

If there are any breaking changes to public APIs, please call them out.
-->
No
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@alamb alamb alamb approved these changes
@Jefffrey Jefffrey Jefffrey approved these changes

Assignees

No one assigned

Labels

arrow Changes to the arrow crate performance

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Json] RunEndEncoded decoder optimization

3 participants

@liamzwbao @alamb @Jefffrey